# /**
# * Copyright (c) 2022 Beijing Jiaotong University
# * PhotLab is licensed under [Open Source License].
# * You can use this software according to the terms and conditions of the [Open Source License].
# *
# * You may obtain a copy of [Open Source License] at: [https://open.source.license/]
# *
# * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# *
# * See the [Open Source License] for more details.
# */
# /**
# * Author: Shihao Xie
# * Created: Jul. 19, 2022
# * Supported by: National Key Research and Development Program of China
# */
import os

from phot_server.core.exception import FilePathError, TemporaryResourceException, Grid2dPlotException
from phot_server.factory.grid_factory import GridFactory
from phot_server.repository.grid_repository import GridRepository
from phot_server.repository.material_repository import MaterialRepository
from phot_server.repository.project_config_repository import ProjectConfigRepository
from phot_server.rpc.device_data_type.ttypes import PhotGrid, PhotWaveguide, PhotSource, \
    PhotFigure, PhotDetector, PhotMaterial, PhotMaterialProperty
from phot_server.server_config.global_config import TEMP_ROOT, FIGURE_POSTFIX, GRID_POSTFIX
from phot_server.utils.generate_uuid import generate_uuid
from phot_server.utils.phtd import parse_phtd, save_phtd
from photfdtd.solve import Solve


class GridManager:
    def __init__(self):
        self.__repository = GridRepository()
        self.__factory = GridFactory()
        self.__project_repository = ProjectConfigRepository()
        self.__material_repository = MaterialRepository()

    # Grid Manage
    def open_grid(self, file_path):
        if file_path is None:
            raise FilePathError("Graph path is None.")
        file_path = os.path.abspath(file_path)
        grid = parse_phtd(file_path)
        grid.gridPath = file_path
        self.__repository.add_grid(grid)
        return self.get_grid(grid.gridId)

    def save_grid(self, grid_id):
        grid = self.__repository.get_grid(grid_id)
        config = self.__project_repository.get_project_config_by_name(grid.name + GRID_POSTFIX)
        if grid.gridPath is not None:
            path = grid.gridPath
        elif config is not None:
            path = os.path.abspath(os.path.join(config.projectDirectory, config.mainFileName + GRID_POSTFIX))
        else:
            raise TemporaryResourceException(grid_id)
        save_phtd(path, grid)

    def save_grid_as(self, grid_id, file_path):
        grid = self.__repository.get_grid(grid_id)
        save_phtd(file_path, grid)
        if grid.gridPath is None:
            grid.gridPath = file_path

    def close_grid(self, grid_id):
        grid = self.get_grid(grid_id)
        self.__repository.delete_grid(grid_id)
        return grid

    def get_opened_grid_list(self):
        grids = self.__repository.get_grids()
        id_name_dict = dict()
        for grid in grids.values():
            id_name_dict.update({str(grid.gridId): str(grid.name)})
        return id_name_dict

    def get_grid(self, grid_id: int):
        return self.__repository.get_grid(grid_id)

    def has_grid(self, grid_id):
        return self.__repository.has_grid(grid_id)

    def add_grid(self, grid: PhotGrid):
        if grid.gridId is None or grid.gridId == "":
            grid.gridId = generate_uuid()
        self.__repository.add_grid(grid)
        return grid

    def set_grid(self, grid: PhotGrid):
        # old_grid = self.__repository.get_grid(grid.gridId)
        # self.__repository.set_grid(grid)
        # for key, value in grid.__dict__.items():
        #     if value is not None and type(value) != list and value != getattr(old_grid, key):
        #         setattr(old_grid, key, value)
        self.update_grid_attrs(grid.gridId, grid.__dict__)

    def update_grid_attrs(self, grid_id, updated_parameters: dict):
        self.__repository.update_grid_attrs(grid_id, updated_parameters)


    def delete_grid(self, grid_id: int):
        self.__repository.delete_grid(grid_id)

    # Light source manage
    def add_source(self, grid_id: int, source: PhotSource):
        if source.sourceId is None:
            source.sourceId = generate_uuid()
        self.__repository.add_source(grid_id, source)
        return source

    def get_source(self, grid_id: int, source_id: int):
        return self.__repository.get_source(grid_id, source_id)

    def update_source_attrs(self, grid_id: int, source: PhotSource):
        self.__repository.set_source(grid_id, source)

    def delete_source(self, grid_id: int, source_id: int):
        self.__repository.delete_source(grid_id, source_id)

    # Waveguide manage
    def add_waveguide(self, grid_id: int, waveguide: PhotWaveguide):
        if waveguide.waveguideId is None:
            waveguide.waveguideId = generate_uuid()
        self.__repository.add_waveguide(grid_id, waveguide)
        return waveguide

    def get_waveguide(self, grid_id: int, waveguide_id: int):
        return self.__repository.get_waveguide(grid_id, waveguide_id)

    def update_waveguide_attrs(self, grid_id: int, waveguide: PhotWaveguide):
        self.__repository.set_waveguide(grid_id, waveguide)

    def delete_waveguide(self, grid_id: int, waveguide_id: int):
        self.__repository.delete_waveguide(grid_id, waveguide_id)

    # detector manage
    def add_detector(self, grid_id: int, detector: PhotDetector):
        if detector.detectorId is None:
            detector.detectorId = generate_uuid()
        self.__repository.add_detector(grid_id, detector)
        return detector

    def get_detector(self, grid_id: int, detector_id: int):
        return self.__repository.get_detector(grid_id, detector_id)

    def update_detector_attrs(self, grid_id: int, detector: PhotDetector):
        self.__repository.set_detector(grid_id, detector)

    def delete_detector(self, grid_id: int, detector_id: int):
        self.__repository.delete_detector(grid_id, detector_id)

    def get_figure(self, grid_id: int, axis: str = "x", axis_index=-1, figure_type: str = "Sectional View",
                   run=False) -> PhotFigure:
        figure_name = "Grid_" + str(grid_id) + "_" + figure_type.replace(" ", "_")

        if run:
            self.fdtd_simulation(grid_id)
        target = self.save_figure_as(grid_id, TEMP_ROOT, axis, axis_index=axis_index, figure_type=figure_type)
        with open(target.path, "rb") as figure:
            figure_binary = figure.read()
        return PhotFigure(name=target.name, axis=target.axis, type=target.type, path=target.path, figure=figure_binary)

    def get_figures(self, grid_id: int, axis_index=None, figure_type="Sectional View") -> dict:
        if axis_index is None:
            axis_index = [0, 0, 0]
        figures = {}
        figure_list = self.save_figures_as(grid_id, TEMP_ROOT, ["x", "y", "z"], axis_index, figure_type)
        for figure in figure_list:
            with open(figure.path, "rb") as f:
                figure_binary = f.read()
                figures[figure.type.split(" ")[0] + "_" + figure.axis] = PhotFigure(name=figure.name, axis=figure.axis,
                                                                                    type=figure.type, path=figure.path,
                                                                                    figure=figure_binary)
        # if figure_type == "Sectional View" or "":
        #     figures.update(
        #         {"Sectional_x": self.get_figure(grid_id, axis="x", axis_index=int(axis_index[0]),
        #                                         figure_type="Sectional View"),
        #          "Sectional_y": self.get_figure(grid_id, axis="y", axis_index=int(axis_index[1]),
        #                                         figure_type="Sectional View"),
        #          "Sectional_z": self.get_figure(grid_id, axis="z", axis_index=int(axis_index[2]),
        #                                         figure_type="Sectional View")})
        # if figure_type == "Refractive Index View" or "":
        #     figures.update(
        #         {"Refractive_x": self.get_figure(grid_id, axis="x", axis_index=int(axis_index[0]),
        #                                          figure_type="Refractive Index View"),
        #          "Refractive_y": self.get_figure(grid_id, axis="y", axis_index=int(axis_index[1]),
        #                                          figure_type="Refractive Index View"),
        #          "Refractive_z": self.get_figure(grid_id, axis="z", axis_index=int(axis_index[2]),
        #                                          figure_type="Refractive Index View")})
        # if figure_type == "Field View" or "":
        #     figures.update(
        #         {"Field_x": self.get_figure(grid_id, axis="x", axis_index=int(axis_index[0]),
        #                                     figure_type="Field View"),
        #          "Field_y": self.get_figure(grid_id, axis="y", axis_index=int(axis_index[1]),
        #                                     figure_type="Field View"),
        #          "Field_z": self.get_figure(grid_id, axis="z", axis_index=int(axis_index[2]),
        #                                     figure_type="Field View")})
        return figures

    def fdtd_simulation(self, grid_id):
        grid = self.get_grid(grid_id)
        fdtd_grid = self.__factory.create_grid_object(grid)
        fdtd_grid.run(animate=False,
                      step=5,
                      axis="z",
                      axis_number=0,
                      time=int(grid.totalTime))
        setattr(fdtd_grid, "isRun", True)
        self.__repository.add_fdtd_grid(grid_id, fdtd_grid)
        # if grid.zLength > 1:
        #     self.save_figure_as(grid_id, os.path.join(TEMP_ROOT, "Grid_" + str(grid_id) + ".png"))
        # else:
        #     self.save_figure_as(grid_id, os.path.join(TEMP_ROOT, "Grid_" + str(grid_id) + "z.png"), axis="z")

    def save_Sectional_View(self, grid_id, axis, axis_index, folder):
        fdtd_grid = self.__repository.get_fdtd_grid(grid_id)
        grid = self.get_grid(grid_id)
        if fdtd_grid is None:
            fdtd_grid = self.__factory.create_grid_object(grid)
            self.__repository.add_fdtd_grid(grid_id, fdtd_grid)
            fdtd_grid = self.__repository.get_fdtd_grid(grid_id)
        # parameter_dict = {axis: int(axis_index),
        #                   "filePath": file_path,
        #                   "save": True,
        #                   "show": False}
        if not os.path.exists(folder):
            os.makedirs(folder, 0o777)
        fdtd_grid.folder = folder
        fdtd_grid.save_fig(axis=axis,
                           axis_number=axis_index,
                           animate=False,
                           time=grid.totalTime,
                           geo=None,
                           show_structure=True,
                           show_energy=True)
        time = grid.totalTime
        if not getattr(fdtd_grid, "isRun", False):
            time = 0
        index = "_%s=%d, total_time=%d" % (axis, axis_index, time)
        return PhotFigure(name="Grid_" + str(grid_id) + "_Sectional_View", type="Sectional View", axis=axis,
                          path=os.path.join(folder, f"file{str(index).zfill(4)}.png"))

    def save_Refractive_Index_View(self, grid_id, axis, axis_index, folder):
        fdtd_grid = self.__repository.get_fdtd_grid(grid_id)
        grid = self.get_grid(grid_id)
        if fdtd_grid is None:
            fdtd_grid = self.__factory.create_grid_object(grid)
            self.__repository.add_fdtd_grid(grid_id, fdtd_grid)
            fdtd_grid = self.__repository.get_fdtd_grid(grid_id)
        if not os.path.exists(folder):
            os.makedirs(folder, 0o777)
        Solve(fdtd_grid, axis=axis, index=int(axis_index), filepath=folder).plot()
        return PhotFigure(name="Grid_" + str(grid_id) + "_Refractive_Index_View", type="Refractive Index View",
                          axis=axis,
                          path=os.path.join(folder, '%s_%s=%d.png' % ('index', axis, axis_index)))

    def save_Field_View(self, grid_id, axis="z", axis_index=0, field="E", folder=""):
        import photfdtd.grid
        fdtd_grid = self.__repository.get_fdtd_grid(grid_id)
        grid = self.get_grid(grid_id)
        if fdtd_grid is None:
            fdtd_grid = self.__factory.create_grid_object(grid)
            self.__repository.add_fdtd_grid(grid_id, fdtd_grid)
            fdtd_grid = self.__repository.get_fdtd_grid(grid_id)
        if not os.path.exists(folder):
            os.makedirs(folder, 0o777)
        photfdtd.grid.Grid.plot_field(grid=fdtd_grid, axis=axis, axis_index=axis_index, field=field, field_axis=None,
                                      folder=folder)
        title = "%s intensity" % field
        return PhotFigure(name="Grid_" + str(grid_id) + "_Field_View", type="Field View", axis=axis,
                          path="%s//%s_%s=%i.png" % (folder, title, axis, axis_index))

    def _create_figure_new_path(self, file_path, grid_id, axis, axis_index, figure_type: str):
        if os.path.isfile(file_path):
            (dir_and_file_name, postfix) = os.path.splitext(file_path)
            path = dir_and_file_name + "_" + "Grid_" + str(grid_id) + "_" + str(axis) + "=" + str(axis_index) + postfix
            return path
        elif os.path.isdir(file_path):
            path = os.path.join(file_path, "Grid_" + str(grid_id), figure_type.replace(" ", "_"))
            return path
        else:
            return file_path

    def save_figure_as(self, grid_id, file_path: str, axis: str = None, axis_index: int = None,
                       figure_type="Sectional View"):
        import matplotlib, shutil
        matplotlib.pyplot.set_loglevel("info")

        grid = self.get_grid(grid_id)
        axis_list = ["x", "y", "z"]
        axis_list.remove(axis)
        for ax in axis_list:
            if getattr(grid, ax + "Length") is None or getattr(grid, ax + "Length") <= 1:
                raise Grid2dPlotException(grid, axis)

        if axis_index == -1:
            axis_index = self.get_mid_index(grid, axis)
        if figure_type == "Sectional View":
            figure_path = self._create_figure_new_path(TEMP_ROOT, grid_id, axis, axis_index, "Sectional View")
            sectional_view_figure = self.save_Sectional_View(grid_id, axis, axis_index, figure_path)
            if not os.path.isdir(file_path):
                shutil.copy(sectional_view_figure.path, file_path)
                sectional_view_figure.path = file_path
            self.__repository.add_fdtd_figure_log(grid_id, [sectional_view_figure])
            return sectional_view_figure
        elif figure_type == "Refractive Index View":
            figure_path = self._create_figure_new_path(TEMP_ROOT, grid_id, axis, axis_index,
                                                       "Refractive Index View")
            refractive_index_view_figure = self.save_Refractive_Index_View(grid_id, axis, axis_index, figure_path)
            if not os.path.isdir(file_path):
                shutil.copy(refractive_index_view_figure.path, file_path)
                refractive_index_view_figure.path = file_path
            self.__repository.add_fdtd_figure_log(grid_id, [refractive_index_view_figure])
            return refractive_index_view_figure
        elif figure_type == "Field View":
            figure_path = self._create_figure_new_path(TEMP_ROOT, grid_id, axis, axis_index, "Field View")
            field_view_figure = self.save_Field_View(grid_id, axis, axis_index, "E", figure_path)
            if not os.path.isdir(file_path):
                shutil.copy(field_view_figure.path, file_path)
                field_view_figure.path = file_path
            self.__repository.add_fdtd_figure_log(grid_id, [field_view_figure])
            return field_view_figure

    def save_figures_as(self, grid_id, folder: str, axis=None, axis_index=None, figure_type="Sectional View"):
        import matplotlib
        matplotlib.pyplot.set_loglevel("info")
        if isinstance(axis, str):
            axis_list = [axis]
        else:
            axis_list = ["x", "y", "z"]
        if axis_index is None:
            axis_index_list = [-1 for i in range(len(axis_list))]
        else:
            axis_index_list = axis_index

        grid = self.get_grid(grid_id)
        for index in range(len(axis_index_list)):
            axis_index_list[index] = self.get_mid_index(grid, axis_list[index])

        figure_log = []
        for index in range(len(axis_list)):
            xyz = ["x", "y", "z"]
            xyz.remove(axis_list[index])
            skip = False
            for ax in xyz:
                if getattr(grid, ax + "Length") is None or getattr(grid, ax + "Length") <= 1:
                    skip = True
            if skip:
                continue

            axis = axis_list[index]
            axis_index = axis_index_list[index]
            if figure_type == "Sectional View" or figure_type == "":
                figure_path = self._create_figure_new_path(folder, grid_id, axis, axis_index, "Sectional View")
                sectional_view_figure = self.save_Sectional_View(grid_id, axis, axis_index, figure_path)
                figure_log.append(sectional_view_figure)
            if figure_type == "Refractive Index View" or figure_type == "":
                figure_path = self._create_figure_new_path(folder, grid_id, axis, axis_index,
                                                           "Refractive Index View")
                refractive_index_view_figure = self.save_Refractive_Index_View(grid_id, axis, axis_index, figure_path)
                figure_log.append(refractive_index_view_figure)
            if figure_type == "Field View" or figure_type == "":
                figure_path = self._create_figure_new_path(folder, grid_id, axis, axis_index, "Field View")
                refractive_index_view_figure = self.save_Field_View(grid_id, axis, axis_index, "E", figure_path)
                figure_log.append(refractive_index_view_figure)
        self.__repository.add_fdtd_figure_log(grid_id, figure_log)
        return figure_log

    def get_mid_index(self, grid: PhotGrid, axis):
        import numpy as np
        if grid.waveguides is None or len(grid.waveguides) == 0:
            return 0
        return int(np.mean([getattr(waveguide, axis) for waveguide in grid.waveguides]).item())

    def get_grid_figure_history(self, grid_id):
        return self.__repository.get_fdtd_figure_logs(grid_id)

    def get_material_name_list(self):
        return self.__material_repository.get_material_name_list()

    def get_material(self, material_name):
        material = PhotMaterial(name=material_name)
        material.color = self.__material_repository.get_material_color(material_name)
        material.properties = [PhotMaterialProperty(propertyId=pid, waveLength=waveLength, real=re, imag=im)
                               for pid, waveLength, re, im in self.__material_repository.
                               get_material_refractive_index(material_name=material_name)]
        return material

    def get_all_materials(self):
        materials = []
        for material_name in self.__material_repository.get_material_name_list():
            material = PhotMaterial(name=material_name)
            material.color = self.__material_repository.get_material_color(material_name)
            material.properties = [PhotMaterialProperty(propertyId=pid, waveLength=waveLength, real=re, imag=im)
                                   for pid, waveLength, re, im in self.__material_repository.
                                   get_material_refractive_index(material_name=material_name)]
            materials.append(material)
        return materials

    def add_material(self, material):
        self.__material_repository.add_material(material_name=material.name,
                                                color=material.color,
                                                refractive_index_points=[(m_property.waveLength,
                                                                          m_property.real, m_property.imag)
                                                                         for m_property in material.properties])
        return material

    def update_material(self, material):
        self.__material_repository.update_material(material_name=material.name,
                                                   color=material.color,
                                                   refractive_index_points=[
                                                       (m_property.propertyId, m_property.waveLength,
                                                        m_property.real, m_property.imag)
                                                       for m_property in material.properties])
        return material

    def update_material_property(self, material_name, material_property):
        self.__material_repository.update_material_property(material_name=material_name,
                                                            refractive_index=material_property)
        return self.get_material(material_name)

    def delete_material(self, material_name):
        self.__material_repository.delete_material(material_name)

    def delete_material_property(self, material_name, property_id):
        self.__material_repository.delete_material_property(material_name, property_id)

    def save_material_as(self, material_name, file_path):
        from phot_server.utils.toml_parser import save_toml
        material_dict = {"name": material_name,
                         "color": self.__material_repository.get_material_color(material_name),
                         "refractive_index": [(waveLength, re, im)
                                              for pid, waveLength, re, im in self.__material_repository.
                                              get_material_refractive_index(material_name=material_name)]}
        material_dict = {"Material": material_dict}
        save_toml(file_path, material_dict)

    def calculate_mode(self, grid_id, axis, index, lam=1.55, neff=3.47, neigs=1, x_boundary_low="pml",
                       y_boundary_low="pml",
                       x_thickness_low=0,
                       y_thickness_low=0, x_boundary_high="zero", y_boundary_high="zero", x_thickness_high=0,
                       y_thickness_high=0):
        grid = self.get_grid(grid_id)
        fdtd_grid = self.__factory.create_grid_object(grid)
        solver = Solve(fdtd_grid)
        if axis == 'x':
            setattr(solver, "n", solver.geometry[index, :, :, :])
        elif axis == 'y':
            setattr(solver, "n", solver.geometry[:, index, :, :])
        elif axis == 'z':
            setattr(solver, "n", solver.geometry[:, :, index, :])
        else:
            raise RuntimeError('axis should be x, y or z! ')
        return solver.calculate_mode(lam, neff, neigs, x_boundary_low, y_boundary_low, x_thickness_low,
                                     y_thickness_low, x_boundary_high, y_boundary_high, x_thickness_high,
                                     y_thickness_high)

    def get_mode_figure(self, grid_id, axis, axis_index, lam=1.55, neff=3.47, neigs=1):
        grid = self.get_grid(grid_id)
        # fdtd_grid = self.__factory.create_grid_object(grid)
        # solver = Solve(fdtd_grid)

        # solver.plot(axis, axis_index, TEMP_ROOT)
        mode_dict = self.calculate_mode(grid_id, axis, axis_index, lam, neff, neigs,
                                        x_boundary_low="pml", y_boundary_low="pml",
                                        x_thickness_low=grid.pmlWidthX, y_thickness_low=grid.pmlWidthY,
                                        x_boundary_high="zero", y_boundary_high=grid.pmlWidthY,
                                        x_thickness_high=grid.pmlWidthX, y_thickness_high=0)
        Solve.draw_mode(TEMP_ROOT, mode_dict, content="amplitude")

        name = '%s_%s=%d.png' % ('index', axis, axis_index)
        path = '%s\\%s_%s=%d.png' % (TEMP_ROOT, 'index', axis, axis_index)
        with open(TEMP_ROOT, "rb") as figure:
            figure_binary = figure.read()
        figure_log = PhotFigure(name=name, axis=axis, type="", path=path)
        self.__repository.add_fdtd_figure_log(grid_id, [figure_log])
        return PhotFigure(name=name, axis=axis, type="", path=path, figure=figure_binary)

    # def calculate_TEfraction(self, grid_id, n_levels=6):
    #     grid = self.get_grid(grid_id)
    #     fdtd_grid = self.__factory.create_grid_object(grid)
    #     solver = Solve(fdtd_grid)
    #     solver.filepath = TEMP_ROOT
    #     solver.calculate_TEfraction(n_levels)

    def get_TEfraction_figure(self, grid_id, n_levels=6):
        grid = self.get_grid(grid_id)
        fdtd_grid = self.__factory.create_grid_object(grid)
        solver = Solve(fdtd_grid)
        solver.filepath = TEMP_ROOT
        solver.calculate_TEfraction(n_levels)

    def sweep(self, grid_id, steps=5, lams=[]):
        grid = self.get_grid(grid_id)
        fdtd_grid = self.__factory.create_grid_object(grid)
        solver = Solve(fdtd_grid)
        solver.filepath = TEMP_ROOT
        solver.sweep(steps, lams)
        name = "wavelength_beta"
        path = '%s\\%s.png' % (TEMP_ROOT, "wavelength_beta")
        with open(TEMP_ROOT, "rb") as figure:
            figure_binary = figure.read()
        figure_log = PhotFigure(name=name, axis="", type="", path=path)
        self.__repository.add_fdtd_figure_log(grid_id, figure_log)
        return PhotFigure(name=name, axis="", type="", path=path, figure=figure_binary)

    def get_waveguide_parameter_list(self, name):
        import photfdtd
        import inspect
        import json
        from inspect import signature

        target_class = [fdtd_cls for cls_name, fdtd_cls in inspect.getmembers(photfdtd, inspect.isclass) if cls_name == name].pop()
        init_signature = signature(target_class.__init__)
        # 创建一个字典来存储参数名和默认值
        params_dict = {}
        # 遍历__init__方法的参数
        for name, parameter in init_signature.parameters.items():
            if name not in ["self", "grid", "reset_xyz", "material"]:
                # 获取参数的默认值，如果没有默认值，则使用None
                default_value = parameter.default if parameter.default is not parameter.empty else None
                # 将参数名和默认值添加到字典中
                params_dict[name] = default_value
        return json.dumps(params_dict)

    def get_all_waveguide_parameter_list(self):
        import photfdtd
        import inspect
        import json
        cls_dict = {}
        for cls_name, fdtd_cls in inspect.getmembers(photfdtd, inspect.isclass):
            if fdtd_cls.__bases__[0].__name__ == "Waveguide" or fdtd_cls.__name__ == "Waveguide":
                cls_dict.update({cls_name: self.get_waveguide_parameter_list(cls_name)})
        return json.dumps(cls_dict)

